#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI营销方案生成系统 - 聊天消息模型
提供AI对话消息的数据模型和数据访问对象

主要功能：
- 聊天消息数据模型
- 消息类型管理
- 消息数据访问对象
- 消息统计分析

设计原则：
- 消息完整性：保证消息数据完整
- 类型区分：区分用户消息和AI回复
- 流式支持：支持流式消息处理
- 性能优化：合理使用索引

Author: AI Marketing System
Date: 2024
"""

from typing import Dict, Any, List, Optional
from datetime import datetime, timezone, timedelta
from pymongo.database import Database
from pymongo import ASCENDING, DESCENDING, TEXT
from enum import Enum

from .base import BaseModel, BaseDAO


class MessageType(Enum):
    """
    消息类型枚举
    """
    USER = 'user'           # 用户消息
    ASSISTANT = 'assistant' # AI助手消息
    SYSTEM = 'system'       # 系统消息
    ERROR = 'error'         # 错误消息


class MessageStatus(Enum):
    """
    消息状态枚举
    """
    SENDING = 'sending'     # 发送中
    SENT = 'sent'          # 已发送
    RECEIVED = 'received'   # 已接收
    PROCESSING = 'processing' # 处理中
    COMPLETED = 'completed' # 已完成
    FAILED = 'failed'      # 失败


class ChatMessage(BaseModel):
    """
    聊天消息模型
    
    字段说明：
    - session_id: 会话ID
    - message_type: 消息类型（user/assistant/system/error）
    - content: 消息内容
    - status: 消息状态
    - sequence: 消息序号（在会话中的顺序）
    - parent_message_id: 父消息ID（用于消息链）
    - tokens: token数量
    - cost: 成本
    - processing_time: 处理时间（秒）
    - model_name: 使用的模型名称
    - model_params: 模型参数
    - raw_response: 原始响应数据
    - error_code: 错误代码
    - error_message: 错误信息
    - metadata: 元数据
    """
    
    collection_name = 'chat_messages'
    
    required_fields = [
        'session_id',
        'message_type',
        'content',
        'status',
        'sequence'
    ]
    
    optional_fields = [
        'parent_message_id',
        'tokens',
        'cost',
        'processing_time',
        'model_name',
        'model_params',
        'raw_response',
        'error_code',
        'error_message',
        'metadata'
    ]
    
    default_values = {
        'status': MessageStatus.SENT.value,
        'sequence': 0,
        'tokens': 0,
        'cost': 0.0,
        'processing_time': 0.0,
        'metadata': {}
    }
    
    def _custom_validate(self) -> List[str]:
        """
        自定义验证逻辑
        
        Returns:
            List[str]: 错误列表
        """
        errors = []
        
        # 验证会话ID
        if hasattr(self, 'session_id') and self.session_id:
            from bson import ObjectId
            try:
                if isinstance(self.session_id, str):
                    ObjectId(self.session_id)
            except Exception:
                errors.append("会话ID格式无效")
        
        # 验证消息类型
        if hasattr(self, 'message_type') and self.message_type:
            valid_types = [msg_type.value for msg_type in MessageType]
            if self.message_type not in valid_types:
                errors.append(f"无效的消息类型: {self.message_type}")
        
        # 验证消息内容
        if hasattr(self, 'content') and self.content is not None:
            if not isinstance(self.content, str):
                errors.append("消息内容必须是字符串")
            if len(self.content.strip()) == 0:
                errors.append("消息内容不能为空")
            if len(self.content) > 100000:  # 100KB限制
                errors.append("消息内容长度不能超过100KB")
        
        # 验证消息状态
        if hasattr(self, 'status') and self.status:
            valid_statuses = [status.value for status in MessageStatus]
            if self.status not in valid_statuses:
                errors.append(f"无效的消息状态: {self.status}")
        
        # 验证序号
        if hasattr(self, 'sequence') and self.sequence is not None:
            if not isinstance(self.sequence, int) or self.sequence < 0:
                errors.append("消息序号必须是非负整数")
        
        # 验证父消息ID
        if hasattr(self, 'parent_message_id') and self.parent_message_id:
            from bson import ObjectId
            try:
                if isinstance(self.parent_message_id, str):
                    ObjectId(self.parent_message_id)
            except Exception:
                errors.append("父消息ID格式无效")
        
        # 验证token数量
        if hasattr(self, 'tokens') and self.tokens is not None:
            if not isinstance(self.tokens, int) or self.tokens < 0:
                errors.append("token数量必须是非负整数")
        
        # 验证成本
        if hasattr(self, 'cost') and self.cost is not None:
            if not isinstance(self.cost, (int, float)) or self.cost < 0:
                errors.append("成本必须是非负数")
        
        # 验证处理时间
        if hasattr(self, 'processing_time') and self.processing_time is not None:
            if not isinstance(self.processing_time, (int, float)) or self.processing_time < 0:
                errors.append("处理时间必须是非负数")
        
        # 验证模型参数
        if hasattr(self, 'model_params') and self.model_params is not None:
            if not isinstance(self.model_params, dict):
                errors.append("模型参数必须是字典格式")
        
        # 验证原始响应
        if hasattr(self, 'raw_response') and self.raw_response is not None:
            if not isinstance(self.raw_response, dict):
                errors.append("原始响应必须是字典格式")
        
        # 验证元数据
        if hasattr(self, 'metadata') and self.metadata is not None:
            if not isinstance(self.metadata, dict):
                errors.append("元数据必须是字典格式")
        
        return errors
    
    def start_processing(self):
        """
        开始处理消息
        """
        self.status = MessageStatus.PROCESSING.value
        self.update_timestamp()
    
    def complete_processing(self, tokens: int = None, cost: float = None, processing_time: float = None):
        """
        完成处理消息
        
        Args:
            tokens: token数量
            cost: 成本
            processing_time: 处理时间
        """
        self.status = MessageStatus.COMPLETED.value
        
        if tokens is not None:
            self.tokens = tokens
        
        if cost is not None:
            self.cost = cost
        
        if processing_time is not None:
            self.processing_time = processing_time
        
        self.update_timestamp()
    
    def fail_processing(self, error_code: str = None, error_message: str = None):
        """
        处理失败
        
        Args:
            error_code: 错误代码
            error_message: 错误信息
        """
        self.status = MessageStatus.FAILED.value
        
        if error_code:
            self.error_code = error_code
        
        if error_message:
            self.error_message = error_message
        
        self.update_timestamp()
    
    def is_user_message(self) -> bool:
        """
        是否是用户消息
        
        Returns:
            bool: 是否是用户消息
        """
        return self.message_type == MessageType.USER.value
    
    def is_assistant_message(self) -> bool:
        """
        是否是AI助手消息
        
        Returns:
            bool: 是否是AI助手消息
        """
        return self.message_type == MessageType.ASSISTANT.value
    
    def is_system_message(self) -> bool:
        """
        是否是系统消息
        
        Returns:
            bool: 是否是系统消息
        """
        return self.message_type == MessageType.SYSTEM.value
    
    def is_error_message(self) -> bool:
        """
        是否是错误消息
        
        Returns:
            bool: 是否是错误消息
        """
        return self.message_type == MessageType.ERROR.value
    
    def is_processing(self) -> bool:
        """
        是否处理中
        
        Returns:
            bool: 是否处理中
        """
        return self.status == MessageStatus.PROCESSING.value
    
    def is_completed(self) -> bool:
        """
        是否已完成
        
        Returns:
            bool: 是否已完成
        """
        return self.status == MessageStatus.COMPLETED.value
    
    def is_failed(self) -> bool:
        """
        是否失败
        
        Returns:
            bool: 是否失败
        """
        return self.status == MessageStatus.FAILED.value
    
    def set_model_info(self, model_name: str, model_params: Dict[str, Any] = None):
        """
        设置模型信息
        
        Args:
            model_name: 模型名称
            model_params: 模型参数
        """
        self.model_name = model_name
        
        if model_params:
            self.model_params = model_params
        
        self.update_timestamp()
    
    def set_raw_response(self, raw_response: Dict[str, Any]):
        """
        设置原始响应
        
        Args:
            raw_response: 原始响应数据
        """
        self.raw_response = raw_response
        self.update_timestamp()
    
    def add_metadata(self, key: str, value: Any):
        """
        添加元数据
        
        Args:
            key: 键
            value: 值
        """
        if not hasattr(self, 'metadata') or self.metadata is None:
            self.metadata = {}
        
        self.metadata[key] = value
        self.update_timestamp()
    
    def get_metadata(self, key: str, default=None):
        """
        获取元数据
        
        Args:
            key: 键
            default: 默认值
            
        Returns:
            元数据值
        """
        if not hasattr(self, 'metadata') or not self.metadata:
            return default
        
        return self.metadata.get(key, default)
    
    def get_content_length(self) -> int:
        """
        获取内容长度
        
        Returns:
            int: 内容长度
        """
        if not hasattr(self, 'content') or not self.content:
            return 0
        
        return len(self.content)
    
    def get_content_preview(self, max_length: int = 100) -> str:
        """
        获取内容预览
        
        Args:
            max_length: 最大长度
            
        Returns:
            str: 内容预览
        """
        if not hasattr(self, 'content') or not self.content:
            return ""
        
        if len(self.content) <= max_length:
            return self.content
        
        return self.content[:max_length] + "..."
    
    @classmethod
    def create_indexes(cls, db: Database):
        """
        创建索引
        
        Args:
            db: 数据库实例
        """
        collection = db[cls.collection_name]
        
        # 创建索引
        indexes = [
            # 会话ID索引
            [('session_id', ASCENDING)],
            # 消息类型索引
            [('message_type', ASCENDING)],
            # 消息状态索引
            [('status', ASCENDING)],
            # 序号索引
            [('sequence', ASCENDING)],
            # 父消息ID索引
            [('parent_message_id', ASCENDING)],
            # 模型名称索引
            [('model_name', ASCENDING)],
            # 创建时间索引
            [('created_at', DESCENDING)],
            # 更新时间索引
            [('updated_at', DESCENDING)],
            # token数量索引
            [('tokens', DESCENDING)],
            # 成本索引
            [('cost', DESCENDING)],
            # 处理时间索引
            [('processing_time', DESCENDING)],
            # 复合索引：会话ID+序号
            [('session_id', ASCENDING), ('sequence', ASCENDING)],
            # 复合索引：会话ID+消息类型
            [('session_id', ASCENDING), ('message_type', ASCENDING)],
            # 复合索引：会话ID+状态
            [('session_id', ASCENDING), ('status', ASCENDING)],
            # 复合索引：消息类型+状态
            [('message_type', ASCENDING), ('status', ASCENDING)],
            # 复合索引：会话ID+创建时间
            [('session_id', ASCENDING), ('created_at', ASCENDING)]
        ]
        
        for index in indexes:
            try:
                collection.create_index(index)
            except Exception as e:
                print(f"创建索引失败 {index}: {e}")


class ChatMessageDAO(BaseDAO):
    """
    聊天消息数据访问对象
    
    提供消息相关的数据库操作：
    - 消息CRUD操作
    - 消息查询和筛选
    - 消息统计和分析
    """
    
    @property
    def collection_name(self) -> str:
        return 'chat_messages'
    
    @property
    def model_class(self):
        return ChatMessage
    
    def find_by_session_id(self, session_id: str, limit: int = None) -> List[ChatMessage]:
        """
        根据会话ID查找消息
        
        Args:
            session_id: 会话ID
            limit: 限制数量
            
        Returns:
            List[ChatMessage]: 消息列表
        """
        return self.find_many(
            filter_dict={'session_id': session_id},
            sort=[('sequence', ASCENDING)],
            limit=limit
        )
    
    def find_by_message_type(self, message_type: str, limit: int = None) -> List[ChatMessage]:
        """
        根据消息类型查找消息
        
        Args:
            message_type: 消息类型
            limit: 限制数量
            
        Returns:
            List[ChatMessage]: 消息列表
        """
        return self.find_many(
            filter_dict={'message_type': message_type},
            sort=[('created_at', DESCENDING)],
            limit=limit
        )
    
    def find_by_status(self, status: str, limit: int = None) -> List[ChatMessage]:
        """
        根据状态查找消息
        
        Args:
            status: 消息状态
            limit: 限制数量
            
        Returns:
            List[ChatMessage]: 消息列表
        """
        return self.find_many(
            filter_dict={'status': status},
            sort=[('created_at', DESCENDING)],
            limit=limit
        )
    
    def find_processing_messages(self) -> List[ChatMessage]:
        """
        查找处理中的消息
        
        Returns:
            List[ChatMessage]: 消息列表
        """
        return self.find_many(
            filter_dict={'status': MessageStatus.PROCESSING.value},
            sort=[('created_at', ASCENDING)]
        )
    
    def find_failed_messages(self, limit: int = 100) -> List[ChatMessage]:
        """
        查找失败的消息
        
        Args:
            limit: 限制数量
            
        Returns:
            List[ChatMessage]: 消息列表
        """
        return self.find_many(
            filter_dict={'status': MessageStatus.FAILED.value},
            sort=[('created_at', DESCENDING)],
            limit=limit
        )
    
    def get_session_message_count(self, session_id: str) -> int:
        """
        获取会话消息数量
        
        Args:
            session_id: 会话ID
            
        Returns:
            int: 消息数量
        """
        return self.collection.count_documents({'session_id': session_id})
    
    def get_next_sequence(self, session_id: str) -> int:
        """
        获取下一个序号
        
        Args:
            session_id: 会话ID
            
        Returns:
            int: 下一个序号
        """
        result = self.collection.find_one(
            {'session_id': session_id},
            sort=[('sequence', DESCENDING)]
        )
        
        if result:
            return result.get('sequence', 0) + 1
        
        return 1
    
    def get_session_conversation(self, session_id: str) -> List[Dict[str, Any]]:
        """
        获取会话对话内容（简化格式）
        
        Args:
            session_id: 会话ID
            
        Returns:
            List[Dict[str, Any]]: 对话列表
        """
        messages = self.find_by_session_id(session_id)
        
        conversation = []
        for message in messages:
            conversation.append({
                'role': message.message_type,
                'content': message.content,
                'sequence': message.sequence,
                'created_at': message.created_at
            })
        
        return conversation
    
    def get_messages_by_date_range(
        self,
        start_date: datetime,
        end_date: datetime
    ) -> List[ChatMessage]:
        """
        根据日期范围查找消息
        
        Args:
            start_date: 开始日期
            end_date: 结束日期
            
        Returns:
            List[ChatMessage]: 消息列表
        """
        return self.find_many(
            filter_dict={
                'created_at': {
                    '$gte': start_date,
                    '$lte': end_date
                }
            },
            sort=[('created_at', DESCENDING)]
        )
    
    def get_message_stats(self) -> Dict[str, Any]:
        """
        获取消息统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        pipeline = [
            {
                '$group': {
                    '_id': {
                        'message_type': '$message_type',
                        'status': '$status'
                    },
                    'count': {'$sum': 1},
                    'avg_tokens': {'$avg': '$tokens'},
                    'avg_cost': {'$avg': '$cost'},
                    'avg_processing_time': {'$avg': '$processing_time'}
                }
            }
        ]
        
        result = list(self.collection.aggregate(pipeline))
        
        stats = {
            'total': 0,
            'by_type': {},
            'by_status': {},
            'avg_tokens': 0,
            'avg_cost': 0,
            'avg_processing_time': 0
        }
        
        total_tokens = 0
        total_cost = 0
        total_processing_time = 0
        
        for item in result:
            group_info = item['_id']
            message_type = group_info['message_type']
            status = group_info['status']
            count = item['count']
            avg_tokens = item.get('avg_tokens', 0)
            avg_cost = item.get('avg_cost', 0)
            avg_processing_time = item.get('avg_processing_time', 0)
            
            stats['total'] += count
            
            # 按类型统计
            if message_type not in stats['by_type']:
                stats['by_type'][message_type] = {'count': 0, 'avg_tokens': 0, 'avg_cost': 0}
            stats['by_type'][message_type]['count'] += count
            
            # 按状态统计
            if status not in stats['by_status']:
                stats['by_status'][status] = {'count': 0, 'avg_tokens': 0, 'avg_cost': 0}
            stats['by_status'][status]['count'] += count
            
            total_tokens += avg_tokens * count
            total_cost += avg_cost * count
            total_processing_time += avg_processing_time * count
        
        if stats['total'] > 0:
            stats['avg_tokens'] = total_tokens / stats['total']
            stats['avg_cost'] = total_cost / stats['total']
            stats['avg_processing_time'] = total_processing_time / stats['total']
        
        return stats
    
    def get_model_usage_stats(self) -> List[Dict[str, Any]]:
        """
        获取模型使用统计
        
        Returns:
            List[Dict[str, Any]]: 模型使用统计列表
        """
        pipeline = [
            {
                '$match': {
                    'model_name': {'$exists': True, '$ne': None}
                }
            },
            {
                '$group': {
                    '_id': '$model_name',
                    'total_messages': {'$sum': 1},
                    'total_tokens': {'$sum': '$tokens'},
                    'total_cost': {'$sum': '$cost'},
                    'avg_processing_time': {'$avg': '$processing_time'}
                }
            },
            {
                '$sort': {'total_messages': -1}
            }
        ]
        
        result = list(self.collection.aggregate(pipeline))
        
        stats = []
        for item in result:
            stats.append({
                'model_name': item['_id'],
                'total_messages': item['total_messages'],
                'total_tokens': item['total_tokens'],
                'total_cost': item['total_cost'],
                'avg_tokens_per_message': item['total_tokens'] / item['total_messages'] if item['total_messages'] > 0 else 0,
                'avg_cost_per_message': item['total_cost'] / item['total_messages'] if item['total_messages'] > 0 else 0,
                'avg_processing_time': item.get('avg_processing_time', 0)
            })
        
        return stats
    
    def cleanup_old_messages(self, days: int = 30) -> int:
        """
        清理旧的消息记录
        
        Args:
            days: 保留天数
            
        Returns:
            int: 删除的记录数
        """
        cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
        
        result = self.collection.delete_many({
            'created_at': {'$lt': cutoff_date},
            'status': {'$in': [
                MessageStatus.COMPLETED.value,
                MessageStatus.FAILED.value
            ]}
        })
        
        return result.deleted_